'''
Author: 渔火Arcadia  https://github.com/yhArcadia
Date: 2025-04-18 00:13:00
LastEditors: 渔火Arcadia
LastEditTime: 2025-04-18 15:02:59
FilePath: /example/utils/loader.py
Description: 

Copyright (c) 2025 by 渔火Arcadia 1761869682@qq.com, All Rights Reserved. 
'''
from torchvision import transforms
from PIL import Image
import torch


# 训练集加载器，返回一个批次的图片张量
def train_loader():
    # 加载数据集（5张图片，共4个分类）
    image_paths = [
        './dataset/traindata/0.jpg',
        './dataset/traindata/1.jpg',
        './dataset/traindata/2-1.jpg',
        './dataset/traindata/2-2.jpg',
        './dataset/traindata/qq.jpg'
    ]
    # 批量加载和转换图片
    transform = transforms.Compose([
        transforms.ToTensor()  # 将PIL Image或numpy.ndarray转为torch.Tensor
    ])
    images = []
    for path in image_paths:
        img = Image.open(path).convert('RGB')
        images.append(transform(img))
    # 堆叠成一个批次(batchsize为5)的张量
    x = torch.stack(images)
    return x



# 测试集加载器，返回一个图片张量
def test_loader(img_path):
    # 加载图片并转换为张量
    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    img = Image.open(img_path).convert('RGB')  # 确保是3通道
    x = transform(img).unsqueeze(0)  # 添加batch维度（只有1张图，batchsize为1）
    return x



# 标签映射器，返回一个字典，键为标签索引，值为标签名称
def label_loader():
    label_map = {
        0: "这是数字0",
        1: "这是数字1", 
        2: "这是数字2",
        3: "这是qq企鹅脸"
        }
    return label_map




if __name__ == '__main__':
    # x = train_loader()
    # print(x)
    y= test_loader('./dataset/testdata/0000.jpg')
    print(y)